# -*- coding: utf-8 -*-
"""
将数据库中以绝对路径存储的 location 字段批量迁移为相对 ROOT_PATH 的路径。

覆盖范围：
- qa_files.location
- attach_2_file.location

执行前建议备份数据库。
"""

import sys
import os

from loguru import logger

os.environ["SQLALCHEMY_DATABASE_URI"] = "mysql+pymysql://root:CowcMjJIDOahClHdW6@39.106.94.192:3306/yoka-tang"

# 允许脚本独立运行
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from configs import RESOURCES_DIR
from utils.file_utils import to_rel_path
from db.base import SessionLocal
from db.models.files import QAFile
from db.models.attachments import Attach2File
from db.models.files import EmergencyFile, AccidentCaseFile, DualControlFile


APPLY_CHANGES = True  # 仅当为 True 时写回数据库；为 False 时仅打印变更预览


def migrate_table_locations(session, model, field_name: str) -> int:
    """将指定表的指定字段由绝对路径迁移为相对路径。
    返回更新条数
    """
    updated = 0
    # 以用户指定的固定前缀进行切分，仅保留其后的相对部分
    fixed_prefix = "/app/portai-server-dev/scripts"
    fixed_prefix_abs = os.path.abspath(fixed_prefix)

    # 逐批处理，避免一次性加载过多
    batch_size = 500
    offset = 0
    while True:
        rows = (session.query(model)
                .offset(offset)
                .limit(batch_size)
                .all())
        if not rows:
            break

        for row in rows:
            value = getattr(row, field_name, None)
            if not value:
                continue

            # 仅当为绝对路径且位于 ROOT_PATH 下时转换
            # 统一绝对路径形式
            try:
                abs_value = os.path.abspath(value)
            except Exception:
                abs_value = value

            # 优先按照固定前缀进行截断
            if isinstance(abs_value, str) and abs_value.startswith(fixed_prefix_abs):
                suffix = abs_value[len(fixed_prefix_abs):].lstrip("/\\")
                if suffix and suffix != value:
                    new_value = suffix.replace('\\', '/')
                    if APPLY_CHANGES:
                        setattr(row, field_name, new_value)
                    else:
                        logger.info(f"[DRY-RUN] {model.__tablename__}.{field_name}: {value} -> {new_value}")
                    updated += 1
            # 回退：若固定前缀不匹配，则按 RESOURCES_DIR 规则转换
            elif os.path.isabs(abs_value):
                new_rel = to_rel_path(abs_value)
                if new_rel and new_rel != value:
                    if APPLY_CHANGES:
                        setattr(row, field_name, new_rel)
                    else:
                        logger.info(f"[DRY-RUN] {model.__tablename__}.{field_name}: {value} -> {new_rel}")
                    updated += 1

        if APPLY_CHANGES and updated % 100 == 0:
            session.flush()

        offset += batch_size

    if APPLY_CHANGES:
        session.commit()
    return updated


def main():
    logger.info(f"RESOURCES_DIR: {RESOURCES_DIR}")
    logger.info("Fixed prefix for split: /home/yoka-main/resources")
    session = SessionLocal()
    try:
        for File in [QAFile, Attach2File, DualControlFile, AccidentCaseFile, EmergencyFile]:
            logger.info(f"开始迁移 {File.__tablename__}.location -> 相对路径")
            qa_updated = migrate_table_locations(session, File, 'location')
            logger.info(f"{File.__tablename__} 更新条数: {qa_updated}")

        # logger.info("开始迁移 attach_2_file.location -> 相对路径")
        # at_updated = migrate_table_locations(session, Attach2File, 'location')
        # logger.info(f"attach_2_file 更新条数: {at_updated}")
    finally:
        session.close()

    logger.info("迁移完成")


if __name__ == "__main__":
    main()


